import numpy as np
import matplotlib.pyplot as plt
import scipy.io as sio
import os, sys

def load_target_mat(target_file):
    data_dict = sio.loadmat(target_file)

    lh_data = data_dict['lh'].transpose()[0]
    rh_data = data_dict['rh'].transpose()[0]

    return lh_data, rh_data

def cal_corr(src_list, des_list):
    corr = np.corrcoef(np.array(src_list), np.array(des_list))
    return corr[0][1]

def plot_diff(output_dir, n_frames, ori_data, des_data, hemi, src_label, tar_label):
    fig, axs = plt.subplots(1, 1)
    axs.plot(n_frames, ori_data, label = src_label)
    axs.plot(n_frames, des_data, label = tar_label)

    corr = cal_corr(ori_data, des_data)
    corr = round(corr, 4)

    axs.set_title('hemi: ' + hemi + ' [r:' + str(corr) + ']')
    axs.set_xlabel('frames')
    axs.set_ylabel('variability')
    axs.legend()

    fig.tight_layout()
    plt.savefig(output_dir + hemi + '_' + src_label + '_'+ tar_label + '.png', dpi=150)


symbol = "GU"
project_dir = '/mri_projects/ASD_Analysis'
target_gu_file = os.path.join(project_dir, 'Results/Variability/intra_Variability_FS4_' + symbol + '.mat')

symbol = "KKI"
target_kki_file = os.path.join(project_dir, 'Results/Variability/intra_Variability_FS4_' + symbol + '.mat')

target_origin_file = '/home/li/Documents/Trainning/SADTestFCMap/Scripts/Variability_FS4.mat'

lh_origin_data, rh_origin_data = load_target_mat(target_origin_file)
lh_gu_whole_data, rh_gu_whole_data = load_target_mat(target_gu_file)
lh_kki_whole_data, rh_kki_whole_data = load_target_mat(target_kki_file)

n_frames = range(0, len(lh_origin_data))               # white noise 2

tar_dir = os.path.join(project_dir, 'Figures/Variability/')

plot_diff(tar_dir, n_frames, lh_origin_data, lh_gu_whole_data, 'lh', 'std', 'gu')
plot_diff(tar_dir, n_frames, lh_origin_data, lh_kki_whole_data, 'lh', 'std', 'kki')
plot_diff(tar_dir, n_frames, lh_gu_whole_data, lh_kki_whole_data, 'lh', 'gu', 'kki')

plot_diff(tar_dir, n_frames, rh_origin_data, rh_gu_whole_data, 'rh', 'std', 'gu')
plot_diff(tar_dir, n_frames, rh_origin_data, rh_kki_whole_data, 'rh', 'std', 'kki')
plot_diff(tar_dir, n_frames, rh_gu_whole_data, rh_kki_whole_data, 'rh', 'gu', 'kki')

wh_origin_data = []
wh_origin_data.extend(lh_origin_data)
wh_origin_data.extend(rh_origin_data)

gu_whole_data = []
gu_whole_data.extend(lh_gu_whole_data)
gu_whole_data.extend(rh_gu_whole_data)

kki_whole_data = []
kki_whole_data.extend(lh_kki_whole_data)
kki_whole_data.extend(rh_kki_whole_data)

n_frames = range(0, len(wh_origin_data))

plot_diff(tar_dir, n_frames, wh_origin_data, gu_whole_data, 'wh', 'std', 'gu')
plot_diff(tar_dir, n_frames, wh_origin_data, kki_whole_data, 'wh', 'std', 'kki')
plot_diff(tar_dir, n_frames, gu_whole_data, kki_whole_data, 'wh', 'gu', 'kki')
